package org.chen.ctrip.service.impl;// OllamaEmbeddingService.java
import com.fasterxml.jackson.annotation.JsonProperty;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;

@Slf4j
@Service
public class OllamaEmbeddingService {

    @Value("${ollama.embedding.base-url:http://8.219.169.114:11434}")
    private String ollamaBaseUrl;

    @Value("${ollama.embedding.model:bge-m3}")
    private String embeddingModel;

    @Value("${ollama.embedding.cache.enable:true}")
    private boolean cacheEnabled;

    @Value("${ollama.embedding.cache.max-size:1000}")
    private int maxCacheSize;

    private final RestTemplate restTemplate;
    private final ObjectMapper objectMapper;
    
    // 简单的内存缓存，避免重复计算相同文本的向量
    private final Map<String, List<Float>> embeddingCache = new ConcurrentHashMap<>();

    public OllamaEmbeddingService(@Qualifier("ollamaRestTemplate") RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
        this.objectMapper = new ObjectMapper();
    }

    /**
     * 生成文本嵌入向量
     */
    public List<Float> generateEmbedding(String text) {
        if (text == null || text.trim().isEmpty()) {
            log.warn("输入文本为空，返回零向量");
            return generateZeroVector();
        }

        // 清理文本
        String cleanText = cleanText(text);
        String cacheKey = generateCacheKey(cleanText);

        // 检查缓存
        if (cacheEnabled && embeddingCache.containsKey(cacheKey)) {
            log.debug("从缓存中获取嵌入向量: {}", cacheKey);
            return embeddingCache.get(cacheKey);
        }

        try {
            log.debug("调用Ollama生成嵌入向量: {}", cleanText.substring(0, Math.min(50, cleanText.length())));
            
            List<Float> embedding = callOllamaEmbedding(cleanText);
            
            // 存入缓存
            if (cacheEnabled && embedding != null) {
                manageCache();
                embeddingCache.put(cacheKey, embedding);
            }
            
            log.debug("成功生成嵌入向量，维度: {}", embedding != null ? embedding.size() : 0);
            return embedding;

        } catch (Exception e) {
            log.error("调用Ollama嵌入服务失败，文本: {}", cleanText.substring(0, Math.min(50, cleanText.length())), e);
            
            // 失败时返回确定性随机向量作为降级方案
            return generateDeterministicVector(cleanText);
        }
    }

    /**
     * 调用Ollama API生成嵌入向量
     */
    private List<Float> callOllamaEmbedding(String text) {
        try {
            // 构建请求
            OllamaEmbeddingRequest request = new OllamaEmbeddingRequest();
            request.setModel(embeddingModel);
            request.setPrompt(text);

            // 设置请求头
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);

            String requestBody = objectMapper.writeValueAsString(request);
            HttpEntity<String> entity = new HttpEntity<>(requestBody, headers);

            // 发送请求
            String url = ollamaBaseUrl + "/api/embeddings";
            ResponseEntity<OllamaEmbeddingResponse> response = restTemplate.exchange(
                url, HttpMethod.POST, entity, OllamaEmbeddingResponse.class);

            if (response.getStatusCode().is2xxSuccessful() && response.getBody() != null) {
                OllamaEmbeddingResponse responseBody = response.getBody();
                
                if (responseBody.getEmbedding() != null && !responseBody.getEmbedding().isEmpty()) {
                    // 将Double列表转换为Float列表
                    List<Float> floatEmbedding = new ArrayList<>();
                    for (Double value : responseBody.getEmbedding()) {
                        floatEmbedding.add(value.floatValue());
                    }
                    return floatEmbedding;
                } else {
                    log.warn("Ollama返回的嵌入向量为空");
                }
            } else {
                log.error("Ollama API调用失败，状态码: {}", response.getStatusCode());
            }

        } catch (Exception e) {
            log.error("调用Ollama API异常", e);
            throw new RuntimeException(e.getMessage());
        }

        return null;
    }

    /**
     * 清理文本，去除不必要的字符和格式
     */
    private String cleanText(String text) {
        if (text == null || text.isEmpty()) {
            return "";
        }

        try {
            // 先进行文本清理操作
            String cleaned = text
                    .trim()
                    .replaceAll("\\s+", " ")  // 多个空白字符替换为单个空格
                    .replaceAll("[\\n\\r\\t]+", " ");  // 换行符等替换为空格

            // 安全地截取字符串，确保不会越界
            final int maxLength = 8000;
            if (cleaned.length() > maxLength) {
                return cleaned.substring(0, maxLength);
            }

            return cleaned;

        } catch (Exception e) {
            log.error("清理文本时发生异常", e);
            // 降级处理：返回原始文本的安全截取版本
            int safeLength = Math.min(text.length(), 8000);
            return text.substring(0, safeLength);
        }
    }

    /**
     * 生成缓存键
     */
    private String generateCacheKey(String text) {
        return String.valueOf(text.hashCode());
    }

    /**
     * 管理缓存大小
     */
    private void manageCache() {
        if (embeddingCache.size() >= maxCacheSize) {
            // 简单的LRU策略：清理一半的缓存
            int removeCount = maxCacheSize / 2;
            embeddingCache.entrySet().removeIf(entry -> removeCount > 0);
            log.info("清理嵌入向量缓存，当前大小: {}", embeddingCache.size());
        }
    }

    /**
     * 生成零向量（默认维度）
     */
    private List<Float> generateZeroVector() {
        List<Float> vector = new ArrayList<>();
        for (int i = 0; i < 1024; i++) { // bge-m3默认维度是1024
            vector.add(0.0f);
        }
        return vector;
    }

    /**
     * 生成确定性随机向量作为降级方案
     */
    private List<Float> generateDeterministicVector(String text) {
        List<Float> vector = new ArrayList<>();
        java.util.Random random = new java.util.Random(text.hashCode());
        for (int i = 0; i < 1024; i++) {
            vector.add(random.nextFloat() * 2 - 1);
        }
        log.warn("使用确定性随机向量作为降级方案");
        return vector;
    }

    // 请求和响应的数据类
    @Data
    public static class OllamaEmbeddingRequest {
        private String model;
        private String prompt;
    }

    @Data
    public static class OllamaEmbeddingResponse {
        private List<Double> embedding;
        
        @JsonProperty("total_duration")
        private Long totalDuration;
        
        @JsonProperty("load_duration") 
        private Long loadDuration;
        
        @JsonProperty("prompt_eval_count")
        private Integer promptEvalCount;
    }
}